#粒子群算法。来源：本人原创
import numpy as np

ratio = 0.1


def function(xArray=np.zeros((1, 4))):
    sum = 0
    for i in range(3):
        sum += 100 * (xArray[i + 1] - xArray[i] ** 2) ** 2
    for i in range(4):
        sum += (xArray[i] - 1) ** 2
    return sum


def mappingRand(randomMatrix: np.array, min: int, max: int):  # 映射
    return min + randomMatrix * (max - min)


def checkVectorLimit(vec: np.array, min: int, max: int):  # 将矩阵中的数值限制在一定范围内

    outPoints = np.where(vec >= max)
    # print(outPoints)
    try:
        for i in range(len(outPoints[0])):
            vec[outPoints[0][i]] = max
        outPoints = np.where(vec <= min)
        for i in range(len(outPoints[0])):
            vec[outPoints[0][i]] = min
        return vec
    except Exception as e:
        print(e)
        return vec


class particleSwarm():
    def __init__(self, dim, size, limit: np.array):
        self.dim = dim  # 粒子群点的维度
        self.size = size  # 粒子群点的多少
        self.limit = limit
        self.velocityMatrix = np.zeros((self.size, self.dim))
        self.positionMatrix = np.zeros((self.size, self.dim))
        self.valueList = np.zeros((1, self.size))  # 当前步计算出的结果
        self.individualOp = np.zeros((1, self.size)) + 1 * 10 ** 10  # 个体最优解
        self.swarmOp = 10 ** 10  # 群体历史最优解
        self.pbest = np.zeros((self.size, self.dim))
        self.gbest = np.zeros((1, self.dim))
        # self.omega = 0
        self.c0Ini = 1
        self.c0 = 1
        self.c1 = 1
        self.c2 = 2
        self.run()

    def run(self):
        global ratio
        self.initParticles()
        steps = 0
        while (1):
            steps += 1

            # i = input()

            self.calcFittingValue()
            self.refreshOptimized()
            self.refreshPosAndVelo()
            # print(self.valueList)
            #
            # print(self.individualOp)
            self.c0 = self.c0Ini *np.e**(-steps/10000)#这个函数是衰减算法。
            if (steps % 1000 == 0):
                print(self.velocityMatrix)
                print(self.c0)
                # print(self.velocityMatrix)
                # print(self.positionMatrix)
                print(function(self.pbest[99]))
                print(steps, ratio)
                print(self.swarmOp, self.gbest, function(self.gbest[0]))

    def refreshOptimized(self):

        for i, val in enumerate(self.valueList[0]):
            if (val < self.individualOp[0, i]):
                self.individualOp[0, i] = val
                self.pbest[i, :] = self.positionMatrix[i, :]  # 更新这个粒子最佳位置的值

        tmpmin = np.min(self.individualOp)
        arg = np.argmin(self.individualOp)
        if (tmpmin < self.swarmOp):
            self.swarmOp = tmpmin
            self.gbest[0, :] = self.positionMatrix[arg, :]

    def calcFittingValue(self):

        for i, rows in enumerate(self.positionMatrix):
            self.valueList[:, i] = function(rows)

    def initParticles(self):
        for col in range(self.dim):  # 粒子群的不同维度
            self.positionMatrix[:, col] = mappingRand(np.random.random((self.size)), self.limit[col][0],
                                                      self.limit[col][1])
            self.velocityMatrix[:, col] = mappingRand(np.random.random((self.size)), self.limit[col][0] * ratio,
                                                      self.limit[col][1] * ratio)

    def refreshPosAndVelo(self):
        self.positionMatrix = self.velocityMatrix + self.positionMatrix
        for i in range(self.positionMatrix.shape[1]):  # 以列为单位，因为每个列是同一维度的。遍历四个维度。
            self.positionMatrix[:, i] = checkVectorLimit(self.positionMatrix[:, i], self.limit[i, 0], self.limit[i, 1])

        self.velocityMatrix = self.velocityMatrix * self.c0 + self.c1 * np.random.random(
            (self.size, self.dim)) * (self.pbest - self.positionMatrix) + \
                              self.c2 * np.random.random((self.size, self.dim)) * (
                                      self.gbest - self.positionMatrix)
        for i in range(self.velocityMatrix.shape[1]):  # 以列为单位，因为每个列是同一维度的。遍历四个维度。
            self.velocityMatrix[:, i] = checkVectorLimit(self.velocityMatrix[:, i], self.limit[i, 0] * ratio,
                                                         self.limit[i, 1] * ratio)


if __name__ == "__main__":
    particleSwarm(4, 1000, np.array([[-30, 30], [-30, 30], [-30, 30], [-30, 30]]))
